import pandas as pd
import plotly.graph_objs as go
import numpy as np

# Load the data
df = pd.read_csv('./data/energy.csv')

# Convert 'Year' to a 5 year period
df['Year'] = df['Year'].div(5).astype(int).mul(5)

# Group the data by 'Region' and 'Year' (5 years) and sum the 'Value' column
grouped = df.groupby(['Region', 'Year'])['Value'].sum()

# Create a new DataFrame with the total energy consumption by 5 years
total = grouped.groupby('Year').sum()
regions = {}
for region in df['Region'].unique():
    values = []
    for year in df['Year'].unique():
        try:
            # Calculate the percentage of energy consumption by region and 5 year period
            value = grouped[region, year] / total[year] * 100
            values.append(value)
        except:
            # If there is no data for a region and 5 year period, append 0
            values.append(0)
    regions[region] = values

decades = sorted(df['Year'].unique())
widths = total / total.sum() * 100


# Delete the nan values from the dictionary
# Convert dict_keys object to a list
keys_list = list(regions.keys())

# Check if there are at least two elements in the dictionary
if len(keys_list) >= 2:
    # Delete the second element (index 1) from the dictionary
    del regions[keys_list[1]]


fig = go.Figure()
for key in regions:
    fig.add_trace(go.Bar(
        name=key,
        y=regions[key],
        x=np.cumsum(widths)-widths,
        width=widths,
        offset=0,
        customdata=np.transpose([decades, widths*regions[key]]),
        texttemplate="%{y} x %{width} =<br>%{customdata[1]}",
        textposition="inside",
        textangle=0,
        textfont_color="white",
        hovertemplate="<br>".join([
            "Lustrum: %{customdata[0]}",
            "Total consumption: %{width}",
            "Percentage%: %{y}",
            "Area: %{customdata[1]}",
        ])
    ))

fig.update_xaxes(
    tickvals=np.cumsum(widths)-widths/2,
    ticktext= ["%s<br>%d" % (l, w) for l, w in zip(decades, widths)]
)

fig.update_xaxes(range=[0,100])
fig.update_yaxes(range=[0,100])

fig.update_layout(
    title_text="Energy consumption by region and 5 year period",
    barmode="stack",
    uniformtext=dict(mode="hide", minsize=10),
)

fig.show()